/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/**
 * @file amp_update_scale.cpp
 */
#include "amp_update_scale_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
// #include "src/common/op_host/op_tiling/op_log.h"

// #include "runtime_util.h"
// #include "op_util.h"

// 匿名命名空间，替代static
namespace
{
    constexpr uint64_t GROWTH_FACTOR_ID = 0;
    constexpr uint64_t BACKOFF_FACTOR_ID = 1;
    constexpr uint64_t GROWTH_INTERVAL = 2;

    #define OPS_CHECK_NULL_WITH_CONTEXT(context, ptr)        \
    if ((ptr) == nullptr)                                    \
    {                                                        \
        std::printf("nullptr error!");                       \
        return ge::GRAPH_SUCCESS;                            \
    }                                                        \

}

namespace optiling
{
    // 定义tiling类
    class AmpUpdateScaleTiling
    {
    public:
        explicit AmpUpdateScaleTiling(gert::TilingContext *context) : TilingContext(context) {};
        ge::graphStatus Init();
        ge::graphStatus RunKernelTiling();
        void TilingDataPrint() const;

    private:
        AmpUpdateScaleTilingData TilingData;
        gert::TilingContext *TilingContext = nullptr;
        float growthFactor = 0;
        float backoffFactor = 0;
        int32_t growthInterval = 0;
    };

    ge::graphStatus AmpUpdateScaleTiling::Init()
    {
        auto attrs = TilingContext->GetAttrs();
        OPS_CHECK_NULL_WITH_CONTEXT(TilingContext, attrs);

        auto growthFactorPtr = attrs->GetAttrPointer<float>(GROWTH_FACTOR_ID);
        OPS_CHECK_NULL_WITH_CONTEXT(TilingContext, growthFactorPtr);
        auto backoffFactorPtr = attrs->GetAttrPointer<float>(BACKOFF_FACTOR_ID);
        OPS_CHECK_NULL_WITH_CONTEXT(TilingContext, backoffFactorPtr);
        auto growthIntervalPtr = attrs->GetAttrPointer<int32_t>(GROWTH_INTERVAL);
        OPS_CHECK_NULL_WITH_CONTEXT(TilingContext, growthIntervalPtr);

        growthFactor = *growthFactorPtr;
        backoffFactor = *backoffFactorPtr;
        growthInterval = static_cast<int32_t>(*growthIntervalPtr);
        return ge::GRAPH_SUCCESS;
    }

    ge::graphStatus AmpUpdateScaleTiling::RunKernelTiling()
    {
        TilingContext->SetBlockDim(1);
        TilingData.set_growthFactor(growthFactor);
        TilingData.set_backoffFactor(backoffFactor);
        TilingData.set_growthInterval(growthInterval);

        TilingData.SaveToBuffer(TilingContext->GetRawTilingData()->GetData(), TilingContext->GetRawTilingData()->GetCapacity());
        TilingContext->GetRawTilingData()->SetDataSize(TilingData.GetDataSize());
        TilingDataPrint();
        return ge::GRAPH_SUCCESS;
    }

    void AmpUpdateScaleTiling::TilingDataPrint() const
    {
        // OP_LOGI(TilingContext->GetNodeName(), "growthFactor:%f.", growthFactor);
        // OP_LOGI(TilingContext->GetNodeName(), "backoffFactor:%f.", backoffFactor);
        // OP_LOGI(TilingContext->GetNodeName(), "growthInterval:%d.", growthInterval);

        // printf("【MY】NodeName:%s, growthFactor:%f, backoffFactor:%f, growthInterval:%d\n",
        //        TilingContext->GetNodeName(), growthFactor, backoffFactor, growthInterval);
    }

    static ge::graphStatus TilingAmpUpdateScale(gert::TilingContext *context)
    {
        AmpUpdateScaleTiling tilingObject(context);
        tilingObject.Init();
        return tilingObject.RunKernelTiling();
    }

    static ge::graphStatus TilingPrepareForAmpUpdateScale(gert::TilingParseContext *context)
    {
        return ge::GRAPH_SUCCESS;
    }

    struct AmpUpdateScaleCompileInfo
    {
    };
    // 注册tiling函数
    IMPL_OP(AmpUpdateScale)
        .Tiling(TilingAmpUpdateScale)
        .TilingParse<AmpUpdateScaleCompileInfo>(TilingPrepareForAmpUpdateScale);
}

namespace ge
{
    static ge::graphStatus InferDataType4AmpUpdateScale(gert::InferDataTypeContext *context)
    {
        auto currentScaleDtype = context->GetInputDataType(0);
        auto growthTrackerDtype = context->GetInputDataType(1);

        context->SetOutputDataType(0, currentScaleDtype);
        context->SetOutputDataType(1, growthTrackerDtype);
        return ge::GRAPH_SUCCESS;
    }

    static ge::graphStatus InferShape4AmpUpdateScale(gert::InferShapeContext *context)
    {
        gert::Shape *updatedScaleShape = context->GetOutputShape(0);
        OPS_CHECK_NULL_WITH_CONTEXT(context, updatedScaleShape);

        gert::Shape *updatedGrowthTrackerShape = context->GetOutputShape(1);
        OPS_CHECK_NULL_WITH_CONTEXT(context, updatedGrowthTrackerShape);

        updatedScaleShape->SetDimNum(0);
        updatedGrowthTrackerShape->SetDimNum(0);

        updatedScaleShape->AppendDim(1);
        updatedGrowthTrackerShape->AppendDim(1);
        return ge::GRAPH_SUCCESS;
    }
    // 注册shape函数
    IMPL_OP(AmpUpdateScale)
        .InferShape(InferShape4AmpUpdateScale)
        .InferDataType(InferDataType4AmpUpdateScale);
} // namespace ops

namespace ops
{
    class AmpUpdateScale : public OpDef
    {
    public:
        explicit AmpUpdateScale(const char *name) : OpDef(name)
        {
            this->Input("current_scale")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
            this->Input("growth_tracker")
                .ParamType(REQUIRED)
                .DataType({ge::DT_INT32})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
            this->Input("found_inf")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
            this->Output("updated_scale")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
            this->Output("updated_growth_tracker")
                .ParamType(REQUIRED)
                .DataType({ge::DT_INT32})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
            this->Attr("growth_factor").AttrType(REQUIRED).Float();
            this->Attr("backoff_factor").AttrType(REQUIRED).Float();
            this->Attr("growth_interval").AttrType(REQUIRED).Int();

            OpAICoreConfig aicore_config;
            aicore_config.DynamicCompileStaticFlag(true)
                .DynamicFormatFlag(true)
                .DynamicRankSupportFlag(true)
                .DynamicShapeSupportFlag(true);
            this->AICore().AddConfig("ascend910b", aicore_config);
            this->AICore().AddConfig("ascend910_93", aicore_config);
            this->AICore().AddConfig("ascend310p", aicore_config);
        }
    };
    // 注册算子
    OP_ADD(AmpUpdateScale);
}
